import tushare as ts
import pandas as pd
import numpy as np
import os
import warnings
from tqdm import tqdm

warnings.filterwarnings("ignore")

# 配置路径参数
stock_list_path = r"F:\量化投资程序\机器算法准备\综合修改版\酒类行业明细.xlsx"
output_dir = r"F:\上市公司综合数据"
output_file = os.path.join(output_dir, f"酒类行业_杜邦分析_统一数据_{pd.Timestamp.now().strftime('%Y%m%d')}.csv")
year = "2023"
period = year + "1231"

# 初始化Tushare
ts.set_token('gx03013e909f633ecb66722df66b360f070426613316ebf06ecd3482')
pro = ts.pro_api()


def load_stock_list():
    """加载酒类行业股票列表"""
    try:
        df = pd.read_excel(stock_list_path)
        print(f"成功加载{len(df)}支股票")
        return df['ts_code'].unique().tolist()
    except Exception as e:
        print(f"股票列表加载失败: {e}")
        return []


def fetch_financial_data(stock_code):
    """获取单支股票财务数据"""
    try:
        # 利润表字段处理
        income_fields = [
            "ts_code", "end_date", "revenue", "operate_profit",
            "income_tax", "n_income_attr_p"
        ]

        # 动态获取利息支出字段
        interest_col = None
        for field in ["interest_exp", "finance_exp"]:
            try:
                income_df = pro.income_vip(
                    ts_code=stock_code,
                    period=period,
                    fields=",".join(income_fields + [field]))
                interest_col = field
                break
            except Exception:
                continue

        if not interest_col:
            # 处理无利息支出字段的情况
            income_df = pro.income_vip(
                ts_code=stock_code,
                period=period,
                fields=",".join(income_fields))
            income_df["interest_exp"] = 0  # 添加默认字段

        # 获取资产负债表
        balance_df = pro.balancesheet_vip(
            ts_code=stock_code,
            period=period,
            fields="ts_code,end_date,total_assets,total_hldr_eqy_exc_min_int"
        )

        # 合并数据
        merged_df = pd.merge(income_df, balance_df, on=["ts_code", "end_date"])
        return merged_df.drop_duplicates(subset=["end_date"], keep="last")
    except Exception as e:
        print(f"\n{stock_code} 数据获取失败: {str(e)[:50]}")
        return pd.DataFrame()


def convert_to_yiwan(df):
    """单位转换与缺失值处理"""
    convert_cols = [
        "revenue", "operate_profit", "income_tax", "interest_exp",
        "n_income_attr_p", "total_assets", "total_hldr_eqy_exc_min_int"
    ]
    for col in convert_cols:
        if col not in df.columns:
            df[col] = 0.0
        df[col] = np.where(df[col].isnull(), 0, df[col])
        df[col] = (df[col].astype(float) / 1e8).round(4)
    return df


def validate_calculations(df):
    """计算结果验证"""
    # 验证ROE范围
    invalid_roe = df[(df["ROE"] < 0) | (df["ROE"] > 1)]
    if not invalid_roe.empty:
        print(f"发现异常ROE值\n{invalid_roe}")
    return df


def dupont_analysis(df):
    """五因素模型计算"""
    field_map = {
        "revenue": "营业收入(亿元)",
        "operate_profit": "营业利润(亿元)",
        "income_tax": "所得税费用(亿元)",
        "interest_exp": "利息支出(亿元)",
        "n_income_attr_p": "净利润(亿元)",
        "total_assets": "总资产(亿元)",
        "total_hldr_eqy_exc_min_int": "股东权益(亿元)"
    }
    df = df.rename(columns=field_map)

    # 安全计算
    with np.errstate(divide='ignore', invalid='ignore'):
        # 1.税收负担率
        taxable_income = df["营业利润(亿元)"] - df["利息支出(亿元)"]
        taxable_income = np.where(taxable_income <= 0, 1, taxable_income)
        df["1.税收负担率"] = np.where(
            df["营业利润(亿元)"] > 0,
            1 - df["所得税费用(亿元)"] / taxable_income,
            0
        ).round(4)

        # 2.利息负担率
        df["2.利息负担率"] = np.where(
            df["营业利润(亿元)"] != 0,
            (df["营业利润(亿元)"] - df["利息支出(亿元)"]) / df["营业利润(亿元)"],
            0
        ).round(4)

        # 3.营业利润率
        df["3.营业利润率"] = np.where(
            df["营业收入(亿元)"] != 0,
            df["营业利润(亿元)"] / df["营业收入(亿元)"],
            0
        ).round(4)

        # 4.资产周转率
        df["4.资产周转率"] = np.where(
            df["总资产(亿元)"] > 0,
            df["营业收入(亿元)"] / df["总资产(亿元)"],
            0
        ).round(2)

        # 5.权益乘数
        df["5.权益乘数"] = np.where(
            df["股东权益(亿元)"] > 0,
            df["总资产(亿元)"] / df["股东权益(亿元)"],
            1
        ).round(2)

    # 计算ROE
    df["ROE"] = (
            df["1.税收负担率"] *
            df["2.利息负担率"] *
            df["3.营业利润率"] *
            df["4.资产周转率"] *
            df["5.权益乘数"]
    ).round(4)

    return validate_calculations(df)


def batch_analysis():
    """批量处理并统一导出"""
    stock_codes = load_stock_list()
    if not stock_codes:
        return

    all_data = []
    pbar = tqdm(stock_codes, desc="处理进度", unit="stock")

    for code in pbar:
        pbar.set_postfix_str(f"正在处理 {code}")
        try:
            df = fetch_financial_data(code)
            if not df.empty:
                df = convert_to_yiwan(df)
                result_df = dupont_analysis(df)
                all_data.append(result_df)
        except Exception as e:
            print(f"\n{code} 处理异常: {str(e)[:50]}")

    if all_data:
        final_df = pd.concat(all_data, ignore_index=True)
        # 字段排序优化
        column_order = ["ts_code", "end_date", "营业收入(亿元)", "营业利润(亿元)",
                        "所得税费用(亿元)", "利息支出(亿元)", "净利润(亿元)", "总资产(亿元)",
                        "股东权益(亿元)", "1.税收负担率", "2.利息负担率", "3.营业利润率",
                        "4.资产周转率", "5.权益乘数", "ROE"]
        final_df = final_df.reindex(columns=column_order)
        final_df = final_df.sort_values(["ts_code", "end_date"])

        final_df.to_csv(output_file, index=False, encoding="utf_8_sig")
        print(f"\n成功处理 {len(final_df['ts_code'].unique())}/{len(stock_codes)} 支股票")
        print(f"统一数据文件已保存至: {output_file}")
    else:
        print("\n未成功获取任何数据")


if __name__ == "__main__":
    os.makedirs(output_dir, exist_ok=True)
    batch_analysis()





